from NPM import *
import matplotlib.pyplot as plt

class Myloss(nn.Module):
    """
    自定义损失函数，参照NPM的论文
    """
    def __init__(self):
        super(Myloss,self).__init__()

    def forward(self, y_hat, y_1, y_2):
        return torch.dist(y_hat[:,:100],npm.embeds(y_1).float(),p=2)+torch.dist(y_hat[:,100:], npm.embeds(y_2).float(),p=2)

npm = NPM().to(Config.DEVICE)

if torch.cuda.device_count() > 1:
    print("使用{}个GPU训练".format(torch.cuda.device_count()))
    npm = nn.DataParallel(npm)

criterion = Myloss()
optm = torch.optim.Adam(npm.parameters())

trainset = TrainDataSet()
train_dataloder = DataLoader(trainset, batch_size=Config.batch_size, shuffle=True)

testset = TestDataSet()
test_dataloder = DataLoader(testset, batch_size=Config.batch_size, shuffle=False)

epoch = [i for i in range(50)]
accuracy = []

with open('log.txt','w',encoding='utf-8') as f:
    for i in range(50):
        npm.train()
        loss=0
        for i_batch, batch_data in enumerate(train_dataloder):
            x = batch_data['x']
            y_1 = batch_data['y_1']
            y_2 = batch_data['y_2']
            x.to(Config.DEVICE)
            y_hat = npm(x)
            loss = criterion(y_hat, y_1, y_2)
            optm.zero_grad()
            loss.backward()
            optm.step()
            if i_batch==0:
                npm.eval()
                correct = 0
                predict_wordpairs = []
                for j_batch, batch_test_data in enumerate(test_dataloder):
                    x = batch_test_data['x']
                    tmp1 = x
                    y = batch_test_data['y']
                    tmp2 = y
                    label = batch_test_data['label']
                    x = x.to(Config.DEVICE)
                    x = npm(x)
                    y = npm.embeds(y)
                    predict = torch.tanh(torch.norm(x[:, 100:] - y, p=1, dim=1) - torch.norm(x[:, :100] - y, p=1, dim=1))
                    zero = torch.zeros_like(predict)
                    one = torch.ones_like(predict)
                    y_hat = torch.where(predict > 0, one, zero)
                    correct += sum(y_hat == label).item()
                    hyponym_predict = [Word2Vec.id2word[x] for x in torch.mul(y_hat, tmp1).numpy().tolist() if x != 0]
                    hypernym_predict = [Word2Vec.id2word[y] for y in torch.mul(y_hat, tmp2).numpy().tolist() if y != 0]
                    for j in range(len(hyponym_predict)):
                        predict_wordpairs.append(hyponym_predict[j] + ' ' + hypernym_predict[j] + '\n')
                acc = correct / len(testset)
                f.write("第{}轮训练后，损失为:{:.4f}，准确率为:{:.4f}".format(i, loss,  acc)+ '\n')
                f.writelines(predict_wordpairs)
                f.flush()
                accuracy.append(acc)
    torch.save(npm,Config.save_path)

print("最高的准确率为{:.4f}".format(max(accuracy)))
plt.plot(epoch, accuracy)
plt.show()